#!/usr/bin/python

import json
import os,sys
import socket
import exceptions
import traceback
from optparse import OptionParser

def sockmsg(msg,host,port,recv=0,broadcast=False,protocol="udp", timeout=None):
    if protocol.lower() == "tcp": 
        sock=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
    elif protocol.lower() == "udp":
        sock=socket.socket(socket.AF_INET,socket.SOCK_DGRAM)
    else:
        print ( "sockmsg():You have to set -P --protocol to 'udp' or 'tcp'!" )
        return 1
    if( broadcast ):
        sock.setsockopt(socket.SOL_SOCKET,socket.SO_BROADCAST,1)
    if timeout != None:
        sock.settimeout(timeout)
    sa=(host,port)
    if(protocol.lower()=="tcp"):
        sock.connect(sa)
        ret=sock.send(msg)
        if(recv>0):
            ret=sock.recv(recv)
            ss="[TCP]RESPONSE:%s" % (ret)
            print ss
    else:
        ret=sock.sendto(msg,sa)
        if(recv>0):
            (ret,saf)=sock.recvfrom(recv)
            ss="[UDP]RESPONSE:%s" % (ret)
            print ss
    sock.close()
    return ret

def toint(v):
    try:
        return int(v)
    except:
        return 0

def dump(e):
    if (str(e)!=" 0 " and str(e) !="0"):
        print "Exception:[" + str(e) + "]" 
        print ( traceback.format_exc() )
    else:
        print "Unkown Exception!"

def init_udp():
    global options
    parser = OptionParser()
    parser.add_option("-m", "--msg",dest="msg", help="Message to send! ", default='{"cmd":"echo","from":"test.py"}')
    parser.add_option("-t", "--tohost",dest="tohost", help="Host to send! ", default='255.255.255.255')
    parser.add_option("-p", "--toport",dest="toport", help="Send Port", default='12345')
    parser.add_option("-r", "--recv",dest="recv", help="Recv bytes", default='8192')
    parser.add_option("-b", "--broadcast",dest="broadcast",help="Option for broadcast",default='1')
    (options, args) = parser.parse_args()
    return (parser,options,args)
    
def udp():
    (parser,ops,args)=init_udp()
    # def sockmsg(msg,host,port,recv=0,broadcast=False,protocol="udp",timeout=None):
    ret=sockmsg(ops.msg,ops.tohost,toint(ops.toport), toint(ops.recv) ,broadcast=toint(ops.broadcast) )
    print "udp() was called!:%s \n" % (ret)

def init_echo():
    global options
    parser = OptionParser()
    parser.add_option("-p", "--port",dest="port", help="Listen Port", default='12345')
    (options, args) = parser.parse_args()
    return (parser,options,args)

def echo():
    (parser,ops,args)=init_echo()
    sock=socket.socket(socket.AF_INET,socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET,socket.SO_BROADCAST,1)
    sock.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1)
    sock.bind(("",toint(ops.port)))
    while(True):
        data,addr=sock.recvfrom(8192)
        print "RecvUDP [%s/%s] from %s\n" % (data,len(data),str(addr))
        sock.sendto(data,addr)

def test_json():
    obj={};
    obj["name"]="obj"
    obj["abc"]="ABC"
    obj[123]="123"
    obj2={}
    obj2["name"]="obj2"
    obj2["obj"]=obj
    ret=sockmsg(json.dumps(obj2),host="127.0.0.1",port=12345, recv=8192  )
    print "test_json() was called!:%s \n" % (ret)

def helpinfo():
    print "Help Message!"

def func(fn):
    funcs={ "udp" : udp , "echo" : echo ,"json" : test_json}
    if fn in funcs.keys():
        return funcs[fn]
    return helpinfo

if __name__ == "__main__":
    fn="udp"
    if len(sys.argv) > 1:
        fn=sys.argv[1]
    try:
        f=func(fn)
        f()
    except Exception,e:
        dump(e)

